#include "pch.h"
#include "IoControl.h"
#include "Driver.h"

_Use_decl_annotations_
VOID Evt_IddIoDeviceControl(
	_In_ WDFDEVICE Device,
	_In_ WDFREQUEST Request,
	_In_ size_t OutputBufferLength,
	_In_ size_t InputBufferLength,
	_In_ ULONG IoControlCode
) {
	//UNREFERENCED_PARAMETER(Device);
	//UNREFERENCED_PARAMETER(Request);
	UNREFERENCED_PARAMETER(OutputBufferLength);
	UNREFERENCED_PARAMETER(InputBufferLength);
	//UNREFERENCED_PARAMETER(IoControlCode);
	NTSTATUS status;
	PrintfDebugString("Evt_IddIoDeviceControl: %ld\n", IoControlCode);

	switch (IoControlCode) {
	case IOCTL_MONITOR_ARRIVE:
	{
		PVOID pUserOutputBuf;
		size_t cbOutputBuf;

		bool bMonitorCreateSuccess;

		MONITOR_ARRIVE_ARG_OUT MonitorArriveArgOut;

		PrintfDebugString("Trying to add a new monitor...\n");

		pUserOutputBuf = NULL;
		cbOutputBuf = 0;

		auto* pContext = WdfObjectGet_AdapterWdfContext(Device);
		auto* pIndirectAdapter = pContext->pIndirectAdapter;
		if (pIndirectAdapter == nullptr) {
			PrintfDebugString("pIndirectAdapter is null.\n");
		}
		// Since adapter is null, we cannot attach new monitors to it.
		if (pIndirectAdapter->IsAdapterNull()) {
			//pContext->pIndirectAdapter = new indirect_disp::IndirectAdapter(Device);
			PrintfDebugString("pIndirectAdapter->IddCxAdapterObject is null! Ignoring this I/O request.\n");
			WdfRequestComplete(Request, STATUS_IO_DEVICE_ERROR);
			return;
		}

		bMonitorCreateSuccess = pIndirectAdapter->NewMonitorArrives(&MonitorArriveArgOut);
		if (!bMonitorCreateSuccess) {
			PrintfDebugString("ARRIVE pIndirectAdapter->NewMonitorArrives Failed!\n");
			WdfRequestComplete(Request, STATUS_IO_DEVICE_ERROR);
			return;
		}

		status = WdfRequestRetrieveOutputBuffer(Request, sizeof(MONITOR_ARRIVE_ARG_OUT),(PVOID*)&pUserOutputBuf, &cbOutputBuf);
		if (!NT_SUCCESS(status)) {
			PrintfDebugString("ARRIVE WdfRequestRetrieveOutputBuffer Failed!\n");
			WdfRequestComplete(Request, status);
			return;
		}

		_Analysis_assume_(cbOutputBuf >= sizeof(MONITOR_ARRIVE_ARG_OUT));
		//Destination<-source
		CopyMemory(pUserOutputBuf, &MonitorArriveArgOut, sizeof(MONITOR_ARRIVE_ARG_OUT));

		WdfRequestCompleteWithInformation(Request, STATUS_SUCCESS, sizeof(MONITOR_ARRIVE_ARG_OUT));
		break;
	}
	case IOCTL_MONITOR_DEPART: {
		PrintfDebugString("DEPART Trying to delete a monitor...\n");
		PVOID pUserInputBuf;
		size_t cbInputBuf;

		bool bMonitorDepartSuccess;

		MONITOR_DEPART_ARG_IN MonitorDepartArgIn;

		pUserInputBuf = NULL;
		cbInputBuf = 0;

		status = WdfRequestRetrieveInputBuffer(Request, sizeof(MONITOR_DEPART_ARG_IN), (PVOID*)&pUserInputBuf, &cbInputBuf);
		if (!NT_SUCCESS(status)) {
			PrintfDebugString("WdfRequestRetrieveInputBuffer Failed!\n");
			WdfRequestComplete(Request, status);
			return;
		}

		_Analysis_assume_(cbInputBuf >= sizeof(MONITOR_DEPART_ARG_IN));
		//Destination<-source
		CopyMemory(&MonitorDepartArgIn, pUserInputBuf, sizeof(MONITOR_DEPART_ARG_IN));

		auto* pContext = WdfObjectGet_AdapterWdfContext(Device);
		auto* pIndirectAdapter = pContext->pIndirectAdapter;

		if (pIndirectAdapter == nullptr) {
			PrintfDebugString("pIndirectAdapter is null.\n");
		}
		// Since adapter is null, we cannot attach new monitors to it.
		if (pIndirectAdapter->IsAdapterNull()) {
			//pContext->pIndirectAdapter = new indirect_disp::IndirectAdapter(Device);
			PrintfDebugString("pIndirectAdapter->IddCxAdapterObject is null! Ignoring this I/O request.\n");
			WdfRequestComplete(Request, STATUS_IO_DEVICE_ERROR);
			return;
		}

		bMonitorDepartSuccess = pIndirectAdapter->DelMonitor(MonitorDepartArgIn);
		if (!bMonitorDepartSuccess) {
			PrintfDebugString("DEPART pIndirectAdapter->MonitorDepart Failed!\n");
			WdfRequestComplete(Request, STATUS_IO_DEVICE_ERROR);
			return;
		}
		WdfRequestComplete(Request, STATUS_SUCCESS);
		break;
	
		
		//bool bMonitorDepartSuccess;
		//MONITOR_DEPART_ARG_IN MonitorDepartArgIn;
		//MonitorDepartArgIn.dwMonitorIndex = 0;
		//
		///*if (1 != pIndirectDeviceCtx->idxFlag) {
		//	WdfRequestComplete(Request, STATUS_SUCCESS);
		//	return;
		//}*/
		//bMonitorDepartSuccess = pIndirectAdapter->DelMonitor(MonitorDepartArgIn);
		//if (!bMonitorDepartSuccess) {
		//	PrintfDebugString("IndirectAdapter->DelMonitor Failed!\n");
		//	WdfRequestComplete(Request, STATUS_IO_DEVICE_ERROR);
		//	return;
		//}

		//WdfRequestComplete(Request, STATUS_SUCCESS);
		break;
	}
	case IOCTL_ADAPTER_ECHO: {
		LPWSTR pwUserInputBuf;
		size_t inputLength;

		LPWSTR pwUserOutputBuf;
		size_t outputLength;

		pwUserInputBuf = NULL;
		pwUserOutputBuf = NULL;
		inputLength = 0;
		outputLength = 0;

		status = WdfRequestRetrieveInputBuffer(Request, 0, (PVOID*)&pwUserInputBuf, &inputLength);
		if (!NT_SUCCESS(status)) {
			PrintfDebugString("InWdfRequestRetrieveInputBuffer Failed!\n");
			WdfRequestComplete(Request, status);
		}

		status = WdfRequestRetrieveOutputBuffer(Request, 256, (PVOID*)&pwUserOutputBuf, &outputLength);
		if (!NT_SUCCESS(status)) {
			PrintfDebugString("WdfRequestRetrieveOutputBuffer Failed!\n");
			WdfRequestComplete(Request, status);
		}

		PrintfDebugString("IO Request: OutputBuf = 0x%p, OutputBufLength = %zu\n", pwUserOutputBuf, outputLength);

		StringCbPrintf(pwUserOutputBuf, outputLength, TEXT("Response from UMDF Driver!"));

		PrintfDebugString("Now pUserOutputBuf: %ws\n", pwUserOutputBuf);


		// Must use WdfRequestCompleteWithInformation to set the 'ByteReturn' field!!
		// If you just use WdfRequestComplete, OS thinks you're returning 0 bytes
		// and nothing will be returned to the buffer sent by DeviceIoControl!
		// Why the hell MSDN Documentation doesn't mention that?
		WdfRequestCompleteWithInformation(Request, status, sizeof(TEXT("Response from UMDF Driver!")));

		break;
	}

	default:
		WdfRequestComplete(Request, STATUS_NOT_IMPLEMENTED);
		break;
	}

}
